from pdf_split import pdf_parser
from langchain.embeddings import DashScopeEmbeddings
import os
from langchain.vectorstores import FAISS

class GeneratedVectors:
    def __init__(self):
        self.DASHSCOPE_API_KEY = "sk-2c282b84a1b34473a136e3e34229d062"
        self.embeddings = DashScopeEmbeddings(
                model="text-embedding-v1",  # 使用DashScope的文本嵌入模型
                dashscope_api_key = self.DASHSCOPE_API_KEY
            )
    def create_vectors(self):
        from langchain.vectorstores import FAISS
        # 创建 FAISS 向量数据库
        PdfParser = pdf_parser()
        all_splits = PdfParser.pdf_split()
        vectordb = FAISS.from_documents(
            documents = all_splits, 
            embedding= self.embeddings
        )
        # 如果需要保存到本地文件
        vectordb.save_local("./faiss_index")  # 保存索引到本地
        print("成功保存向量数据库 ")
        # 后续可以从本地加载
        #vectordb = FAISS.load_local("faiss_index", embeddings)
        return vectordb
    
    def get_vectors(self):
        if not os.path.exists ("./faiss_index"):
            vectordb = self.create_vectors()
        else:
            vectordb = FAISS.load_local("faiss_index", self.embeddings)
            print("成功加载向量数据库 ")
        return vectordb


